
exception BadCommandLine of string;;
exception NonzeroExitStatusFailure of int;;
exception SignalFailure of int;;


(* Parses the given command line and returns (program, argv_array),
   where program is the program name and argv_array is an array of
   the argv, the first element of which is the program.  Raises
   BadCommandLine on error. *)
let parse_command_line cmd_line =
  let argv_list = Str.split (Str.regexp " +") cmd_line in
  if List.length argv_list = 0 then
    raise (BadCommandLine "command line is empty")
  else
    let argv_array = Array.of_list argv_list in
    let prog = Array.get argv_array 0 in
    (prog, argv_array);;

(* Creates a subprocess with pipes for stdin, stdout, and stderr.  Returns
   (pid, sub_stdin, sub_stdout, sub_stderr). *)
let create_process_with_pipes prog argv_array =
  let stdout_r,stdout_w = Unix.pipe () in
  let stderr_r,stderr_w = Unix.pipe () in
  let stdin_r,stdin_w = Unix.pipe() in
  let pid = Unix.create_process prog argv_array stdin_r stdout_w stderr_w in
  Unix.close stdout_w;
  Unix.close stderr_w;
  Unix.close stdin_r;
  (pid, (Unix.out_channel_of_descr stdin_w),
        (Unix.in_channel_of_descr stdout_r),
        (Unix.in_channel_of_descr stderr_r));;

(* Tracks the status of a reader thread. *)
type reader_status =
  | Reading
  | ReadSucceeded
  | ReadFailedAndKilled of exn
  | ReadFailedButCouldNotKill of exn;;

(* Tracks the status of the runner thread. *)
type runner_status = Running | RunSucceeded | RunFailed of exn;;

let run_subprocess cmd_line stdout_reader stderr_reader killer_setter =
  (* Split the command_line into an argument array. *)
  let prog, argv_array = parse_command_line cmd_line in

  (* Start the process with pipes. *)
  let pid, sub_stdin, sub_stdout, sub_stderr =
    create_process_with_pipes prog argv_array in
      
  (* Now that we know the pid, we can call the killer_setter with
     a function that will kill that pid if invoked. *)
  killer_setter (fun () ->
                   try Unix.kill pid 9
                   with Unix.Unix_error(Unix.ESRCH, "kill", "") -> ());

  (* Thread that runs a reader for stdout or stderr and handles any
     uncaught exceptions.  On success, sets its status to ReadSucceded.
     On failure, sets its status to ReadFailed with the exception and
     whether or not it managed to kill the process. *)
  let reader_wrapper (reader, status, channel) =
    try
      reader channel;
      status := ReadSucceeded
    with e ->
      begin
        try
          Unix.kill pid 9;
          status := ReadFailedAndKilled e
        with Unix.Unix_error(Unix.ESRCH, _, _) ->
          status := ReadFailedButCouldNotKill e
      end in
      
  (* Start the stdout reader thread, tracking its status. *)
  let stdout_status = ref Reading in
  let stdout_reader_thread =
    Thread.create reader_wrapper (stdout_reader,
                                  stdout_status, sub_stdout) in

  (* Start the stderr reader thread, tracking its status. *)
  let stderr_status = ref Reading in
  let stderr_reader_thread =
    Thread.create reader_wrapper (stderr_reader,
                                  stderr_status, sub_stderr) in
      
  (* Wait for the process to finish. *)
  let _, status = Unix.waitpid [] pid in

  (* Compute the runner's status.  This is just the opinion that
     comes out of waitpid, and we'll reconcile it with the opinions
     of the reader functions in a moment. *)
  let runner_status = match status with
    | Unix.WEXITED(0) -> RunSucceeded
    | Unix.WEXITED(i) -> RunFailed (NonzeroExitStatusFailure i)
    | Unix.WSIGNALED(i) -> RunFailed (SignalFailure i)
    | Unix.WSTOPPED(i) -> RunFailed (SignalFailure i) in
      
  (* Join the reader threads.  This waits for them to finish and
     ensures that stdout_status and stderr_status are populated. *)
  Thread.join stdout_reader_thread;
  Thread.join stderr_reader_thread;

  (* Close our end of the pipes. *)
  close_in sub_stdout;
  close_in sub_stderr;
  close_out sub_stdin;

  (* Now we reconcile the opinions of waitpid with the opinions of
     the readers in order to determine what the actual final status
     is. *)
  match runner_status with
    | Running -> failwith "Shouldn't still be running."
    | RunSucceeded ->
        begin
          (match !stdout_status with
             | Reading -> failwith "Shouldn't still be reading."
             | ReadSucceeded -> ()
             | ReadFailedAndKilled e -> raise e
             | ReadFailedButCouldNotKill e -> raise e);
          (match !stderr_status with
             | Reading -> failwith "Shouldn't still be reading."
             | ReadSucceeded -> ()
             | ReadFailedAndKilled e -> raise e
             | ReadFailedButCouldNotKill e -> raise e);
        end
    | RunFailed (SignalFailure _) ->
        begin
          (match !stdout_status with
             | Reading -> failwith "Shouldn't still be reading."
             | ReadSucceeded -> ()
             | ReadFailedAndKilled e2 -> raise e2
             | ReadFailedButCouldNotKill _ -> ());
          (match !stderr_status with
             | Reading -> failwith "Shouldn't still be reading."
             | ReadSucceeded -> ()
             | ReadFailedAndKilled e2 -> raise e2
             | ReadFailedButCouldNotKill _ -> ());
          raise (SignalFailure 9);
        end
    | RunFailed e -> raise e;
